import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
import matplotlib.colors as mcolors
import random
from collections import Counter
import numpy as np

import config

"""
draw one 'timestamp-machine' grid-basd plot
"""
def draw_machines_res(tasks,start_times,freqs,transaction_map):
    set_machines = set(task['fix'] for task in tasks.values())
    n_machines = len(set_machines)
    H = int(np.lcm.reduce([tasks[i]['T'] for i in range(len(tasks))]))
    
    # Set grid parameters
    rows = n_machines
    cols = H  # Number of rows and columns in grid
    cell_size = 1.0    # Size of each grid cell

    # Create figure and axis
    fig, ax = plt.subplots(figsize=(12, 8))
    fig.patch.set_facecolor('#f8f9fa')  # Set background color

    # Set axis limits
    ax.set_xlim(0, cols * cell_size)
    ax.set_ylim(0, rows * cell_size)
    ax.set_aspect('equal')  # Maintain aspect ratio

    # Turn off axis
    ax.set_axis_off()

    # Add title
    plt.title('WTSN scheduling result', fontsize=18, pad=20, color='#2c3e50')

    # Draw grid lines
    for i in range(rows + 1):
        ax.axhline(y=i * cell_size, color="#212d2d", linestyle='-', linewidth=1.5, alpha=0.7)

    for j in range(cols + 1):
        ax.axvline(x=j * cell_size, color='#95a5a6', linestyle='-', linewidth=1.5, alpha=0.7)

    # Get a set of nice colors
    colors = list(mcolors.TABLEAU_COLORS.values())
    random.shuffle(colors)
    colors_selected = colors[0:config.FREQ_NUM]

    # Randomly select grid cells to color and add annotations
    annotated_cells = []
    for k,task in tasks.items():
        for i,s in enumerate(start_times[k]):
            # Select random color
            color = colors_selected[freqs[k][i]]
            
            # map start_time in one hyper_period
            if s >= H:
                s = s%H
            
            # Create rectangle and add to plot
            rect = Rectangle((s * cell_size, (task['fix']-1) * cell_size), 
                            cell_size, cell_size, 
                            facecolor=color, 
                            edgecolor='#34495e',
                            alpha=0.85,
                            linewidth=1.5)
            ax.add_patch(rect)
            
            # Add annotation (random letter)
            letter = transaction_map[task['transaction']]
            ax.text((s + 0.5) * cell_size, ((task['fix']-1) + 0.5) * cell_size, 
                letter, 
                fontsize=10, 
                color='white' if mcolors.rgb_to_hsv(mcolors.to_rgb(color))[2] < 0.6 else 'black',
                ha='center', va='center',
                weight='bold')
            
            # Record annotated cell
            annotated_cells.append((s, (task['fix']-1)))

    # Add legend text
    # legend_text = f"Each colored cell contains a random letter annotation\nGrid size: {rows}×{cols}, Colored cells: {len(annotated_cells)}"
    # plt.figtext(0.5, 0.01, legend_text, ha='center', fontsize=12, color='#7f8c8d')

    # Add row and column labels
    for i in range(rows):
        ax.text(-0.5, (i + 0.5) * cell_size, f'machine {i+1}', 
                ha='right', va='center', fontsize=10, color='#3498db')
        
    for j in range(cols):
        ax.text((j + 0.5) * cell_size, -0.5, f'{j}', 
                ha='center', va='top', fontsize=10, color='#3498db')

    # Add decorative border
    border = Rectangle((0, 0), cols * cell_size, rows * cell_size, 
                    fill=False, edgecolor='#2c3e50', linewidth=3, alpha=0.8)
    ax.add_patch(border)

    # Adjust layout
    plt.tight_layout()
    plt.subplots_adjust(bottom=0.1)

    # Show plot
    plt.show()
    return


"""
draw two 'timestamp-machine' grid-based plots
"""
def draw_subplot_2_1_res(tasks,start_times,freqs,transaction_map):
    set_machines = set()
    for task in tasks.values():
        set_machines.add(task.get('fix'))
        set_machines.update(task.get('next',[]))
    n_machines = len(set_machines)
    H = int(np.lcm.reduce([tasks[i]['T'] for i in range(len(tasks))]))
    
    machines_to_row = {m:idx for idx,m in enumerate(sorted(set_machines))}
    
    # Set grid parameters
    rows = n_machines
    cols = H  # Number of rows and columns in grid
    cell_size = 1.0    # Size of each grid cell

    # Create figure and axis
    fig, (ax1,ax2) = plt.subplots(2,1,figsize=(12, 8))
    fig.patch.set_facecolor('#f8f9fa')  # Set background color
    
    #set titles
    fig.suptitle('WTSN scheduling result', fontsize=18, color='#2c3e50')
    ax1.set_title('FH TX', fontsize=16, pad=10, color='#3498db')
    ax2.set_title('FH RX', fontsize=16, pad=10, color='#3498db')
    
    # Get a set of nice colors
    colors = list(mcolors.TABLEAU_COLORS.values())
    colors_selected = colors[0:config.FREQ_NUM]

    for ax_idx,ax in enumerate([ax1,ax2]):
        # Set axis limits
        ax.set_xlim(0, cols * cell_size)
        ax.set_ylim(0, rows * cell_size)
        ax.set_aspect('equal')  # Maintain aspect ratio

        # Turn off axis
        ax.set_axis_off()

        # Draw grid lines
        for i in range(rows + 1):
            ax.axhline(y=i * cell_size, color="#212d2d", linestyle='-', linewidth=1.5, alpha=0.7)

        for j in range(cols + 1):
            ax.axvline(x=j * cell_size, color='#95a5a6', linestyle='-', linewidth=1.5, alpha=0.7)


        annotated_cells = []
        for k,task in tasks.items():
            for i,s in enumerate(start_times[k]):
                # Select random color
                color = colors_selected[freqs[k][i]]
                
                # map start_time in one hyper_period
                if s >= H:
                    s = s%H
                    
                if ax_idx == 0:
                    mach_idx = task['fix'] - 1
                else:
                    if len(task['succ']) != 0:
                        mach_idx = (tasks[task['succ'][0]])['fix']-1
                    else:
                        continue
                
                # Create rectangle and add to plot
                rect = Rectangle((s * cell_size, mach_idx * cell_size), 
                                cell_size, cell_size, 
                                facecolor=color, 
                                edgecolor='#34495e',
                                alpha=0.85,
                                linewidth=1.5)
                ax.add_patch(rect)
                
                # Add annotation (random letter)
                letter = transaction_map[task['transaction']]
                ax.text((s + 0.5) * cell_size, (mach_idx + 0.5) * cell_size, 
                    letter, 
                    fontsize=10, 
                    color='white' if mcolors.rgb_to_hsv(mcolors.to_rgb(color))[2] < 0.6 else 'black',
                    ha='center', va='center',
                    weight='bold')
                
                # Record annotated cell
                annotated_cells.append((s, mach_idx))

        # Add legend text
        # legend_text = f"Each colored cell contains a random letter annotation\nGrid size: {rows}×{cols}, Colored cells: {len(annotated_cells)}"
        # plt.figtext(0.5, 0.01, legend_text, ha='center', fontsize=12, color='#7f8c8d')

        # Add row and column labels
        for i in range(rows):
            ax.text(-0.5, (i + 0.5) * cell_size, f'machine {i+1}', 
                    ha='right', va='center', fontsize=10, color='#3498db')
            
        for j in range(cols):
            ax.text((j + 0.5) * cell_size, -0.5, f'{j}', 
                    ha='center', va='top', fontsize=10, color='#3498db')

        # Add decorative border
        border = Rectangle((0, 0), cols * cell_size, rows * cell_size, 
                        fill=False, edgecolor='#2c3e50', linewidth=3, alpha=0.8)
        ax.add_patch(border)

    # Adjust layout
    plt.tight_layout()
    plt.subplots_adjust(bottom=0.1)

    # Show plot
    plt.show()
    return


"""
draw "timestamp-machine" grid-based plots combined TX and RX
"""
def draw_combined_machine_timestamp_res(tasks,start_times,freqs,transaction_map):
    # get machines' set
    set_machines = set()
    for task in tasks.values():
        set_machines.add(task.get('fix'))
        set_machines.update(task.get('next',[]))
    n_machines = len(set_machines)
    H = int(np.lcm.reduce([tasks[i]['T'] for i in range(len(tasks))]))

    # Set grid parameters
    rows = n_machines
    cols = H  # Number of rows and columns in grid
    cell_size = 1.0    # Size of each grid cell
        
    # Create figure and axis
    fig, ax = plt.subplots(figsize=(12, 6))
    fig.patch.set_facecolor('#f8f9fa')
    fig.suptitle('Combined WTSN Scheduling (TX & RX)', fontsize=18, color='#2c3e50')

    # Set axis limits
    ax.set_xlim(0, cols * cell_size)
    ax.set_ylim(0, rows * cell_size)
    ax.set_aspect('equal')
    ax.set_axis_off()

    # Draw grid lines
    for i in range(rows + 1):
        ax.axhline(y=i * cell_size, color="#212d2d", linestyle='-', linewidth=1.5, alpha=0.7)
    for j in range(cols + 1):
        ax.axvline(x=j * cell_size, color='#95a5a6', linestyle='-', linewidth=1.5, alpha=0.7)
        
    # Track occupied cells to detect conflicts
    occupied_cells = set()
    colors = list(mcolors.TABLEAU_COLORS.values())
    colors_selected = colors[:config.FREQ_NUM]  # Assuming config.FREQ_NUM is defined        

    # Function to add rectangle with conflict checking
    def add_rectangle(s, mach_idx, freq, letter, is_tx=True):
        nonlocal occupied_cells
        s_mod = s % H
        
        # Check for conflict
        cell_key = (s_mod,freq, mach_idx)
        if cell_key in occupied_cells:
            raise ValueError(f"Grid conflict at time {s_mod}, machine {mach_idx+1},freq {freq}!")
        occupied_cells.add(cell_key)
        
        # Create rectangle
        color = colors_selected[freq]
        rect = Rectangle((s_mod * cell_size, mach_idx * cell_size),
                         cell_size, cell_size,
                         facecolor=color, 
                         edgecolor='#34495e',
                         alpha=0.85,
                         linewidth=1.5,
                         hatch='//' if not is_tx else None)  # Differentiate RX with hatch
        ax.add_patch(rect)
        
        # Add annotation
        ax.text((s_mod + 0.5) * cell_size, (mach_idx + 0.5) * cell_size,
               f"{letter}{'T' if is_tx else 'R'}",
               fontsize=9,
               color='white' if mcolors.rgb_to_hsv(mcolors.to_rgb(color))[2] < 0.6 else 'black',
               ha='center', va='center',
               weight='bold')

    # Draw all tasks
    for k, task in tasks.items():
        for i, s in enumerate(start_times[k]):
            # TX task (fixed machine)
            tx_mach = task['fix'] - 1
            tx_letter = transaction_map[task['transaction']]
            add_rectangle(s, tx_mach, freqs[k][i], tx_letter, is_tx=True)
            
            # RX task (successor machine if exists)
            if 'next' in task and task['next']:
                rx_mach = task['next'][0] - 1
                rx_letter = transaction_map[task['transaction']]
                add_rectangle(s, rx_mach, freqs[k][i], rx_letter, is_tx=False)

    # Add row and column labels
    for i in range(rows):
        ax.text(-0.5, (i + 0.5) * cell_size, f'm{i+1}', 
                ha='right', va='center', fontsize=10, color='#3498db', weight='bold')
    for j in range(cols):
        if j % 5 == 0:  # Label every 5th time slot for clarity
            ax.text((j + 0.5) * cell_size, -0.5, f'{j}', 
                    ha='center', va='top', fontsize=8, color='#3498db')

    # Add time axis label
    ax.text((cols * cell_size) / 2, -1.5, 'Time Slots', 
            ha='center', va='top', fontsize=12, color='#2c3e50')
    
    # Add machine axis label
    ax.text(-1.5, (rows * cell_size) / 2, 'Machines', 
            ha='center', va='center', fontsize=12, color='#2c3e50',
            rotation=90)

    # Add legend outside the plot area (bottom center)
    from matplotlib.patches import Patch
    legend_elements_tasks = [
        Patch(facecolor='gray', label='TX ', alpha=0.85),
        Patch(facecolor='gray', label='RX ', alpha=0.85, hatch='//')
    ]
    
    # Place legend at the bottom center of the figure
    legend1 = ax.legend(handles=legend_elements_tasks, 
              loc='lower center', 
              bbox_to_anchor=(0.9, 1.05),  # Below the plot
              ncol=2, 
              framealpha=0.9,
              fontsize=10)
    
    # 关键：把第一个 legend 添加到 axes 上
    ax.add_artist(legend1)

    # Add legend of FH
    legend_elements_tasks = []
    for i in range(config.FREQ_NUM):
        legend_FH = Patch(facecolor=colors_selected[i], label=f'FH{i} ', alpha=0.85)
        legend_elements_tasks.append(legend_FH)

    ax.legend(handles=legend_elements_tasks, 
            loc='lower center', 
            bbox_to_anchor=(0.1, 1.05),  # Below the plot
            ncol=2, 
            framealpha=0.9,
            fontsize=10)

    # Add decorative border
    border = Rectangle((0, 0), cols * cell_size, rows * cell_size, 
                       fill=False, edgecolor='#2c3e50', linewidth=2)
    ax.add_patch(border)

    plt.tight_layout()
    plt.subplots_adjust(top=0.92)
    plt.show()


"""
draw 'timestamp-freq' grid-based plots per machine
"""
def draw_timestamp_freq_res(tasks,start_times,freqs,transaction_map):
    set_machines = set()
    for task in tasks.values():
        set_machines.add(task.get('fix'))
        set_machines.update(task.get('next',[]))
        
    n_machines = len(set_machines)
    H = int(np.lcm.reduce([tasks[i]['T'] for i in range(len(tasks))]))

    # freq transmission flag
    TX = 1
    RX = 0

    #find machine freq<-> start_time
    # mach_freq_table["machine_id"]={"start_time:(freq_id_tx/rx,transaction_id)"}
    mach_freq_table= {}
    for k,task in tasks.items():
        if task['fix'] not in mach_freq_table:
            mach_freq_table[task['fix']] = []
        if task['next'][0] not in mach_freq_table:
            mach_freq_table[task['next'][0]] = []
        for i,s in enumerate(start_times[k]):
            mach_freq_table[task['fix']].append((freqs[k][i],TX,task['transaction'],s))
            mach_freq_table[task['next'][0]].append((freqs[k][i],RX,task['transaction'],s))
    
    for k,v in mach_freq_table.items():
        print(f"{k}:{v}")
    

    for mach_k,mach in mach_freq_table.items():
        # Set grid parameters
        rows = config.FREQ_NUM # one shows tx_freq, the other shows rx_freq
        cols = H  # Number of rows and columns in grid
        cell_size = 1.0    # Size of each grid cell

        # Create figure and axis
        fig, ax = plt.subplots(figsize=(12, 8))
        fig.patch.set_facecolor('#f8f9fa')  # Set background color

        # Set axis limits
        ax.set_xlim(0, cols * cell_size)
        ax.set_ylim(0, rows * cell_size)
        ax.set_aspect('equal')  # Maintain aspect ratio

        # Turn off axis
        ax.set_axis_off()

        # Add title
        plt.title(f'WTSN timestamp-freq result machine{mach_k}', fontsize=18, pad=20, color='#2c3e50')

        # Draw grid lines
        for i in range(rows + 1):
            ax.axhline(y=i * cell_size, color="#212d2d", linestyle='-', linewidth=1.5, alpha=0.7)

        for j in range(cols + 1):
            ax.axvline(x=j * cell_size, color='#95a5a6', linestyle='-', linewidth=1.5, alpha=0.7)

        # Get a set of nice colors
        colors = list(mcolors.TABLEAU_COLORS.values())
        # colors_selected = colors[:config.FREQ_NUM]
        color = colors[0]

        # Track occupied cells to detect conflicts
        occupied_cells = set()
        # Track occupied cells to detect whether machines send or receive simultaneously on the same slot
        simul_check_cells = set()

        #FH TX
        for t in mach:
            s= t[0]    # freq choice
            isTx = t[1]   #freq tx/rx flag
            trans = t[2]   # transaction id
            k = t[3]  # start_time
        
            # map start_time in one hyper_period
            k = k % H

            # check repetition
            cell_key = (k,s)
            if cell_key in occupied_cells:
                raise ValueError(f"machine {mach_k} Grid conflict at time {k}, freq {s}!")
            occupied_cells.add(cell_key)
            
            simul_key = (k,isTx)
            if simul_key in simul_check_cells:
                raise ValueError(f"machine {mach_k} simultane conflict at time {k}")
            simul_check_cells.add(simul_key)
            
            # Create rectangle and add to plot
            rect = Rectangle((k * cell_size, s * cell_size), 
                            cell_size, cell_size, 
                            facecolor=color, 
                            edgecolor='#34495e',
                            alpha=0.85,
                            linewidth=1.5,
                            hatch='//' if not isTx else None)
            ax.add_patch(rect)
        
            # Add annotation (random letter)
            letter = transaction_map[trans]
            ax.text((k + 0.5) * cell_size, (s + 0.5) * cell_size, 
                letter, 
                fontsize=10, 
                color='white' if mcolors.rgb_to_hsv(mcolors.to_rgb(color))[2] < 0.6 else 'black',
                ha='center', va='center',
                weight='bold')
            
            # Add legend outside the plot area (bottom center)
        from matplotlib.patches import Patch
        legend_elements_tasks = [
            Patch(facecolor='gray', label='TX ', alpha=0.85),
            Patch(facecolor='gray', label='RX ', alpha=0.85, hatch='//')
        ]
        
        # Place legend at the bottom center of the figure
        legend1 = ax.legend(handles=legend_elements_tasks, 
                loc='upper center', 
                bbox_to_anchor=(0.9, 1.65),  # Below the plot
                ncol=2, 
                framealpha=0.9,
                fontsize=10)

        # Add row and column labels
        for i in range(rows):
            ax.text(-0.5, (i + 0.5) * cell_size, f'FH{i}', 
                    ha='right', va='center', fontsize=10, color='#3498db')
            
        for j in range(cols):
            if j % 5==0:
                ax.text((j + 0.5) * cell_size, -0.5, f'{j}', 
                        ha='center', va='top', fontsize=10, color='#3498db')

        # Add decorative border
        border = Rectangle((0, 0), cols * cell_size, rows * cell_size, 
                        fill=False, edgecolor='#2c3e50', linewidth=3, alpha=0.8)
        ax.add_patch(border)

        # Adjust layout
        plt.tight_layout()
        plt.subplots_adjust(top=0.9)

        # Show plot
        plt.show()
